% main_equilibrium
% Jones, Philippon, Venkateswaran (2021)
%

clear 
clc

%%

warning('Off','all') ;

path(pathdef) ;

%% 

planner = 0 ;

e0exog   = 1 ;
Imax     = 0.2;              %% targets for mortality
BaseMort = 0.0025;           
PeakMort = 3*BaseMort;
deterministic_vac = 1 ;
stoch_prob = 77/78 ;

beta       = (0.97)^(1/52)*stoch_prob ; 
ee_T_vac   = zeros(1,1000); 
ee_T_e0    = zeros(1,1000); 
ee_T_d     = zeros(1,1000); 
if deterministic_vac; beta = (0.97)^(1/52); T_vac = 1*52 ; ee_T_vac(T_vac:end)=0.01; for ii=(T_vac+1):length(ee_T_vac); ee_T_vac(ii)=ee_T_vac(ii-1)*1.2; end ; ee_T_vac(ee_T_vac>1) = 1; end
e0         = e0exog        + (1-e0exog)*(1/2) ;
ec1        = e0exog*0.0001 + (1-e0exog)*((1-e0)/2) ;
el1        = e0exog*0.0001 + (1-e0exog)*((1-e0)/2) ;
R0         = 2.0 ;
chibar1    = 100.5 ;
kappa      = 0.15;                                    % Fraction severely sick
rho        = 0.35;                                    % Recovery rate of infected 
deltabar   = BaseMort/(1-BaseMort)*rho/kappa;         % baseline death rate 
phi        = log(1-deltabar+PeakMort/(1-PeakMort)*rho/kappa)/(Imax*kappa);
N          = 1 ;
eta        = 2 ;
rhom       = 1 ;
Deltachi   = 0.34 ;
us         = 0.5 ;
ud         = 2.5 ;
ee_i0_init = (1/100)*N ; ee_T_i = zeros(1,1000) ; ee_T_i(1) = ee_i0_init ;
psim       = 0.99 ;
deltascale = 1 ;

rhof        = 0.0 ;
omega1      = 0.0 ;
omega2      = 2 ;

rhoa       = 0.0 ;
varphi1    = 0.5 ;
varphi2    = 0.2 ;
amean      = 1 ;

if planner
    model_name = 'infections_planner' ;
    set_params_planner
else
    model_name = 'infections' ;
    set_params
end

run_dynare


figure(1)

set(gcf, 'DefaultAxesLineWidth', 1.0);
set(gcf, 'DefaultLineLineWidth', 1.5);
set(gcf, 'DefaultAxesTickLabelInterpreter','latex'); 
set(gcf, 'DefaultLegendInterpreter','latex');
set(gcf, 'DefaultAxesFontSize',5);

subplot(3,3,2); hold on;
h1=plot(2:(length(s)-1),e(3:end)) ;
line([0 300],[1 1],'LineWidth',1,'Color','k') ;
title('Exposure, $e$','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

subplot(3,3,1); 
yyaxis left; hold on ;
h1 = plot(0:75,c1(1:76),'-','Color',h1.Color) ;
title('Cons (Solid), Labor (Dashed)','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;
set(gca,'YColor','k') ;
ylim([0.8 1.05]) ;

yyaxis right; hold on ;
h1 = plot(0:75,l1(1:76),'--','Color',h1.Color) ;
title('Cons (Solid), Labor (Dashed)','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;
set(gca,'YColor','k') ;
ylim([0.8 1.05]) ;

subplot(3,3,3); hold on;
h1=plot(0:(length(s)-1),m1(1:end)) ;
title('Working from Home, $m_t$','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

subplot(3,3,4); hold on;
h1=plot(0:(length(s)-1),mbar1(1:end)) ;
title('Learning, $\bar{m}_t$','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

subplot(3,3,5); hold on;
h1=plot(1:(length(s)-1),lambdae(2:end)) ;
title('$\lambda_{e,t}$','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

subplot(3,3,6); hold on;
h1=plot(1:(length(s)-1),Vs(2:end)-Vi(2:end)) ;
title('$V_{s,t}-V_{i,t}$','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

subplot(3,3,7); hold on;
h1 = plot(0:(length(s)-1),100*in) ;
title('Infected, \%','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

subplot(3,3,8); hold on;
yyaxis left
hold on ;
plot(0:(length(s)-1),100*s,'-','Color',h1.Color) ;
xlim([0 75]); xticks([0:25:75]) ;
set(gca,'YColor','k') ;

yyaxis right
hold on ;
plot(0:(length(d)-1),100*d,'--','Color',h1.Color) ;
xlim([0 75]); xticks([0:25:75]) ;
set(gca,'YColor','k') ;

title('Susceptible (LHS), Dead (RHS), \%','Interpreter','latex')

subplot(3,3,9); hold on;

CaseMortality = kappa*delta./(rho+kappa*delta);

plot(0:(length(l1)-1),100*CaseMortality) ;
title('Infection Mortality, $\frac{\kappa\delta_t}{\gamma+\kappa\delta_t}$, \%','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

set(findall(gcf, 'Type', 'Text'),'FontWeight', 'Normal')

%% 

planner = 0 ;

e0exog   = 0 ;
Imax     = 0.2;              %% targets for mortality
BaseMort = 0.0025;           
PeakMort = 3*BaseMort;
deterministic_vac = 1 ;
stoch_prob = 77/78 ;

beta       = (0.97)^(1/52)*stoch_prob ; 
ee_T_vac   = zeros(1,1000); 
ee_T_e0    = zeros(1,1000); 
ee_T_d     = zeros(1,1000); 
if deterministic_vac; beta = (0.97)^(1/52); T_vac = 1*52 ; ee_T_vac(T_vac:end)=0.01; for ii=(T_vac+1):length(ee_T_vac); ee_T_vac(ii)=ee_T_vac(ii-1)*1.2; end ; ee_T_vac(ee_T_vac>1) = 1; end
e0         = e0exog        + (1-e0exog)*(1/2) ;
ec1        = e0exog*0.0001 + (1-e0exog)*((1-e0)/2) ;
el1        = e0exog*0.0001 + (1-e0exog)*((1-e0)/2) ;
R0         = 2.0 ;
chibar1    = 100.5 ;
kappa      = 0.15;                                    % Fraction severely sick
rho        = 0.35;                                    % Recovery rate of infected 
deltabar   = BaseMort/(1-BaseMort)*rho/kappa;         % baseline death rate 
phi        = log(1-deltabar+PeakMort/(1-PeakMort)*rho/kappa)/(Imax*kappa);
N          = 1 ;
eta        = 2 ;
rhom       = 1 ;
Deltachi   = 0.34 ;
us         = 0.5 ;
ud         = 2.5 ;
ee_i0_init = (1/100)*N ; ee_T_i = zeros(1,1000) ; ee_T_i(1) = ee_i0_init ;
psim       = 0.99 ;
deltascale = 1 ;

rhof        = 0.0 ;
omega1      = 0.0 ;
omega2      = 2 ;

rhoa       = 0.0 ;
varphi1    = 0.5 ;
varphi2    = 0.2 ;
amean      = 1 ;

if planner
    model_name = 'infections_planner' ;
    set_params_planner
else
    model_name = 'infections' ;
    set_params
end

run_dynare


figure(1)

set(gcf, 'DefaultAxesLineWidth', 1.0);
set(gcf, 'DefaultLineLineWidth', 1.5);
set(gcf, 'DefaultAxesTickLabelInterpreter','latex'); 
set(gcf, 'DefaultLegendInterpreter','latex');
set(gcf, 'DefaultAxesFontSize',5);

subplot(3,3,2); hold on;
h1=plot(2:(length(s)-1),e(3:end)) ;
line([0 300],[1 1],'LineWidth',1,'Color','k') ;
title('Exposure, $e$','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

subplot(3,3,1); 
yyaxis left; hold on ;
h1 = plot(0:75,c1(1:76),'-','Color',h1.Color) ;
title('Cons (Solid), Labor (Dashed)','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;
set(gca,'YColor','k') ;
ylim([0.8 1.05]) ;

yyaxis right; hold on ;
h1 = plot(0:75,l1(1:76),'--','Color',h1.Color) ;
title('Cons (Solid), Labor (Dashed)','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;
set(gca,'YColor','k') ;
ylim([0.8 1.05]) ;

subplot(3,3,3); hold on;
h1=plot(0:(length(s)-1),m1(1:end)) ;
title('Working from Home, $m_t$','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

subplot(3,3,4); hold on;
h1=plot(0:(length(s)-1),mbar1(1:end)) ;
title('Learning, $\bar{m}_t$','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

subplot(3,3,5); hold on;
h1=plot(1:(length(s)-1),lambdae(2:end)) ;
title('$\lambda_{e,t}$','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

subplot(3,3,6); hold on;
h1=plot(1:(length(s)-1),Vs(2:end)-Vi(2:end)) ;
title('$V_{s,t}-V_{i,t}$','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

subplot(3,3,7); hold on;
h1 = plot(0:(length(s)-1),100*in) ;
title('Infected, \%','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

subplot(3,3,8); hold on;
yyaxis left
hold on ;
plot(0:(length(s)-1),100*s,'-','Color',h1.Color) ;
xlim([0 75]); xticks([0:25:75]) ;
set(gca,'YColor','k') ;

yyaxis right
hold on ;
plot(0:(length(d)-1),100*d,'--','Color',h1.Color) ;
xlim([0 75]); xticks([0:25:75]) ;
set(gca,'YColor','k') ;

title('Susceptible (LHS), Dead (RHS), \%','Interpreter','latex')

subplot(3,3,9); hold on;

CaseMortality = kappa*delta./(rho+kappa*delta);

plot(0:(length(l1)-1),100*CaseMortality) ;
title('Infection Mortality, $\frac{\kappa\delta_t}{\gamma+\kappa\delta_t}$, \%','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

set(findall(gcf, 'Type', 'Text'),'FontWeight', 'Normal')

%% 

planner = 0 ;

e0exog   = 0 ;
Imax     = 0.2;              %% targets for mortality
BaseMort = 0.0025;           
PeakMort = 3*BaseMort;
deterministic_vac = 1 ;
stoch_prob = 77/78 ;

beta       = (0.97)^(1/52)*stoch_prob ; 
ee_T_vac   = zeros(1,1000); 
ee_T_e0    = zeros(1,1000); 
ee_T_d     = zeros(1,1000); 
if deterministic_vac; beta = (0.97)^(1/52); T_vac = 1*52 ; ee_T_vac(T_vac:end)=0.01; for ii=(T_vac+1):length(ee_T_vac); ee_T_vac(ii)=ee_T_vac(ii-1)*1.2; end ; ee_T_vac(ee_T_vac>1) = 1; end
e0         = e0exog        + (1-e0exog)*(1/2) ;
ec1        = e0exog*0.0001 + (1-e0exog)*((1-e0)/2) ;
el1        = e0exog*0.0001 + (1-e0exog)*((1-e0)/2) ;
R0         = 2.0 ;
chibar1    = 0.5 ;
kappa      = 0.15;                                    % Fraction severely sick
rho        = 0.35;                                    % Recovery rate of infected 
deltabar   = BaseMort/(1-BaseMort)*rho/kappa;         % baseline death rate 
phi        = log(1-deltabar+PeakMort/(1-PeakMort)*rho/kappa)/(Imax*kappa);
N          = 1 ;
eta        = 2 ;
rhom       = 1 ;
Deltachi   = 0.34 ;
us         = 0.5 ;
ud         = 2.5 ;
ee_i0_init = (1/100)*N ; ee_T_i = zeros(1,1000) ; ee_T_i(1) = ee_i0_init ;
psim       = 0.99 ;
deltascale = 1 ;

rhof        = 0.0 ;
omega1      = 0.0 ;
omega2      = 2 ;

rhoa       = 0.0 ;
varphi1    = 0.5 ;
varphi2    = 0.2 ;
amean      = 1 ;

if planner
    model_name = 'infections_planner' ;
    set_params_planner
else
    model_name = 'infections' ;
    set_params
end

run_dynare


figure(1)

set(gcf, 'DefaultAxesLineWidth', 1.0);
set(gcf, 'DefaultLineLineWidth', 1.5);
set(gcf, 'DefaultAxesTickLabelInterpreter','latex'); 
set(gcf, 'DefaultLegendInterpreter','latex');
set(gcf, 'DefaultAxesFontSize',5);

subplot(3,3,2); hold on;
h1=plot(2:(length(s)-1),e(3:end)) ;
line([0 300],[1 1],'LineWidth',1,'Color','k') ;
title('Exposure, $e_t$','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

subplot(3,3,1); 
yyaxis left; hold on ;
h1 = plot(0:75,c1(1:76),'-','Color',h1.Color) ;
title('Cons (Solid), Labor (Dashed)','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;
set(gca,'YColor','k') ;
ylim([0.8 1.05]) ;

yyaxis right; hold on ;
h1 = plot(0:75,l1(1:76),'--','Color',h1.Color) ;
title('Cons (Solid), Labor (Dashed)','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;
set(gca,'YColor','k') ;
ylim([0.8 1.05]) ;

h=legend('Exog','No WFH','WFH','location','southeast','box','off');
set(h,'Interpreter','latex','FontSize',5) ;

subplot(3,3,3); hold on;
h1=plot(0:(length(s)-1),m1(1:end)) ;
title('Working from Home, $m_t$','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

subplot(3,3,4); hold on;
h1=plot(0:(length(s)-1),mbar1(1:end)) ;
title('Learning, $\bar{m}_t$','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

subplot(3,3,5); hold on;
h1=plot(1:(length(s)-1),lambdae(2:end)) ;
title('$\lambda_{e,t}$','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

subplot(3,3,6); hold on;
h1=plot(1:(length(s)-1),Vs(2:end)-Vi(2:end)) ;
title('$V_{s,t}-V_{i,t}$','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

subplot(3,3,7); hold on;
h1 = plot(0:(length(s)-1),100*in) ;
title('Infected, \%','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

subplot(3,3,8); hold on;
yyaxis left
hold on ;
plot(0:(length(s)-1),100*s,'-','Color',h1.Color) ;
xlim([0 75]); xticks([0:25:75]) ;
set(gca,'YColor','k') ;

yyaxis right
hold on ;
plot(0:(length(d)-1),100*d,'--','Color',h1.Color) ;
xlim([0 75]); xticks([0:25:75]) ;
set(gca,'YColor','k') ;

title('Susc. (LHS), Dead (RHS), \%','Interpreter','latex')

subplot(3,3,9); hold on;

CaseMortality = kappa*delta./(rho+kappa*delta);

plot(0:(length(l1)-1),100*CaseMortality) ;
title('Infection Mort., $\frac{\kappa\delta_t}{\gamma+\kappa\delta_t}$, \%','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ;

set(findall(gcf, 'Type', 'Text'),'FontWeight', 'Normal')

print -depsc ./output/equilibrium_detvaccine.eps
print -dpng ./output/equilibrium_detvaccine.png
close
